#include <stdio.h>
#include <cuda.h>
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <time.h>
#include <sys/time.h>

#define THREAD_PER_BLOCK 256
#define NUM_PER_THREAD 4

__device__ void warpReduce(volatile float* cache,int tid){
    cache[tid]+=cache[tid+32];
    cache[tid]+=cache[tid+16];
    cache[tid]+=cache[tid+8];
    cache[tid]+=cache[tid+4];
    cache[tid]+=cache[tid+2];
    cache[tid]+=cache[tid+1];
}

__global__ void reduce1(float *d_in,float *d_out){
    __shared__ float sdata[THREAD_PER_BLOCK];

    // each thread loads one element from global to shared mem
    unsigned int tid = threadIdx.x;
    unsigned int di = blockIdx.x * (blockDim.x * NUM_PER_THREAD) + threadIdx.x;
    float sum = 0;
    for (unsigned int i = 0; i < NUM_PER_THREAD; i++) {
        sum += d_in[di + i * blockDim.x];
    }
    sdata[tid] = sum;
    __syncthreads();

    // do reduction in shared mem
    for(unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    if (tid < 32) warpReduce(sdata, tid);

    /*
    for(unsigned int s = 1; s < blockDim.x; s <<= 1) {
        if ((tid & (s * 2 - 1)) == 0) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }*/

    // write result for this block to global mem
    if (tid == 0) d_out[blockIdx.x] = sdata[0];
}

bool check(float *out,float *res,int n){
    for(int i=0;i<n;i++){
        if(out[i]!=res[i])
            return false;
    }
    return true;
}

int main(){
    const int N = 1024*1024*1024;
    float *a=(float *)malloc(N*sizeof(float));
    float *d_a;
    cudaMalloc((void **)&d_a,N*sizeof(float));

    int block_num = N / NUM_PER_THREAD / THREAD_PER_BLOCK;
    float *out=(float *)malloc((N/THREAD_PER_BLOCK)*sizeof(float));
    float *d_out;
    cudaMalloc((void **)&d_out,(N/THREAD_PER_BLOCK)*sizeof(float));
    float *res=(float *)malloc((N/THREAD_PER_BLOCK)*sizeof(float));

    for(int i=0;i<N;i++){
        a[i]=1;
    }

    for(int i=0;i<block_num;i++){
        float cur=0;
        for(int j=0;j<THREAD_PER_BLOCK;j++){
            for (int k = 0; k < NUM_PER_THREAD; k++) {
                cur += a[i*THREAD_PER_BLOCK * NUM_PER_THREAD + j * NUM_PER_THREAD + k];
            }
        }
        res[i]=cur;
    }

    cudaMemcpy(d_a,a,N*sizeof(float),cudaMemcpyHostToDevice);

    cudaEvent_t start, end;
    cudaEventCreate(&start);
    cudaEventRecord(start);

    dim3 Grid(N / NUM_PER_THREAD / THREAD_PER_BLOCK, 1);
    dim3 Block(THREAD_PER_BLOCK, 1);

    reduce1<<<Grid,Block>>>(d_a,d_out);

    cudaEventCreate(&end);
    cudaEventRecord(end);

    cudaEventSynchronize(end);

    float msec;
    cudaEventElapsedTime(&msec, start, end);
    printf("GPU time: %3f ms\n", (double)(msec) / CLOCKS_PER_SEC * 1000.0);

    cudaMemcpy(out,d_out,block_num*sizeof(float),cudaMemcpyDeviceToHost);

    if(check(out,res,block_num))printf("the ans is right, %lf, %d\n", out[0], block_num);
    else{
        printf("the ans is wrong, %lf, %d\n", out[0], block_num);
        printf("\n");
    }

    cudaFree(d_a);
    cudaFree(d_out);
}
